import torch
import torch.nn as nn
import torchvision.transforms as transforms
from torchvision.datasets import Caltech101
from torchvision import models
from torch.utils.data import DataLoader, random_split
import matplotlib.pyplot as plt
from torchvision import datasets

# 加载预训练的 ConvNeXt-Tiny 模型
# model = models.convnext_tiny(pretrained=True)
model = models.convnext_tiny(weights=None)

# 数据预处理和增强
transform_train = transforms.Compose([
    transforms.RandomResizedCrop((224, 224)),  # 随机裁剪并调整大小到 224x224
    transforms.RandomHorizontalFlip(),  # 随机水平翻转
    transforms.ToTensor(),  # 转换为张量
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),  # 归一化
])

transform_val = transforms.Compose([
    transforms.Resize((224, 224)),  # 调整大小到 224x224
    transforms.ToTensor(),  # 转换为张量
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),  # 归一化
])

data_dir = '/root/dr/CI/data_tiny_imagenet/tiny-imagenet-200/'
# 加载训练集
train_dataset = datasets.ImageFolder(root=data_dir+'train/', transform=transform_train)
train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True, num_workers=4)

# 加载验证集
val_dataset = datasets.ImageFolder(root=data_dir+'val/', transform=transform_val)
val_loader = DataLoader(val_dataset, batch_size=128, shuffle=False, num_workers=4)

# 替换最后的分类层以适应 tiny_imagenet 的 200 个类别
 # ConvNeXt-Tiny 的最后一层输出是 768 维
model.classifier[2] = torch.nn.Linear(in_features=768, out_features=200)

# 将模型移到 GPU（如果可用）
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
print(f'Training on device: {device}')

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
# 记录训练过程中的损失和精度
train_losses = []
train_accuracies = []
val_accuracies = []

# 定义计算准确率的函数
def calculate_accuracy(loader, model):
    correct = 0
    total = 0
    model.eval()  # 切换到评估模式
    with torch.no_grad():
        for inputs, labels in loader:
            inputs, labels = inputs.to(device), labels.to(device)
            outputs = model(inputs)
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    return correct / total

num_epochs = 100 # 训练 100 个 epoch
# 训练模型（简单示例）
for epoch in range(num_epochs):
    if epoch == 20:
        for param_group in optimizer.param_groups:
            param_group['lr'] = 0.005
        print(f"Learning rate adjusted to {optimizer.param_groups[0]['lr']} at epoch {epoch}")
    if epoch == 60:
        for param_group in optimizer.param_groups:
            param_group['lr'] = 0.0001
        print(f"Learning rate adjusted to {optimizer.param_groups[0]['lr']} at epoch {epoch}")
    if epoch == 80:
        for param_group in optimizer.param_groups:
            param_group['lr'] = 0.00002
        print(f"Learning rate adjusted to {optimizer.param_groups[0]['lr']} at epoch {epoch}")
    model.train()  # 切换到训练模式
    running_loss = 0.0
    correct = 0
    total = 0
    for inputs, labels in train_loader:
        inputs, labels = inputs.to(device), labels.to(device)

        # 前向传播
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        
        # 反向传播和优化
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

        # 计算训练准确率
        _, predicted = torch.max(outputs, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

    train_loss = running_loss / len(train_loader)
    train_acc = correct / total
    val_acc = calculate_accuracy(val_loader, model)

    train_losses.append(train_loss)
    train_accuracies.append(train_acc)
    val_accuracies.append(val_acc)

    print(f"Epoch {epoch+1}/{num_epochs}, Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.4f}, Val Acc: {val_acc:.4f}")


# 绘制训练损失、训练精度和验证精度的曲线图
epochs_range = range(1, num_epochs + 1)
plt.figure(figsize=(10, 5))

# 绘制训练损失
plt.subplot(1, 2, 1)
plt.plot(epochs_range, train_losses, label='Train Loss')
plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.title('Train Loss')
plt.legend()

# 绘制训练精度和验证精度
plt.subplot(1, 2, 2)
plt.plot(epochs_range, train_accuracies, label='Train Accuracy')
plt.plot(epochs_range, val_accuracies, label='Val Accuracy')
plt.xlabel('Epochs')
plt.ylabel('Accuracy')
plt.title('Train and Val Accuracy')
plt.legend()

# 保存图像到工作目录
plt.tight_layout()
plt.savefig('training_results.png')
plt.show()

# 保存模型的 state_dict 到磁盘
torch.save(model.state_dict(), 'convnext_tinyimagenet_100epochs.pth')

print("模型已保存为 convnext_imagenet_100epochs.pth")